import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.pipeline import Pipeline
import warnings
warnings.filterwarnings('ignore')

data = pd.read_csv('data.csv', encoding='gbk')
print(data.head())

# 为了可以更好的进行可视化效果，只取前两个特征进行操作
x = data.iloc[:, 3:5]
y = data.iloc[:, 2:3]
print(x[:10])
# 设置标签，女为0，男为1
lable = LabelEncoder()
y = lable.fit_transform(y)
print(y[:10])

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3)
# 管道方式进行操作
model = Pipeline([
        # 决策树使用ID3算法进行处理，深度为3
        ('DTC', DecisionTreeClassifier(criterion='entropy', max_depth=3))])
model = model.fit(x_train, y_train)
y_test_hat = model.predict(x_test)
print(y_test_hat)

N, M = 100, 100  # 横纵各采样多少个值
plt.rcParams['font.sans-serif'] = ['SimHei']
x1_min, x1_max = x.iloc[:, 0].min(), x.iloc[:, 0].max()  # 第0列的范围
x2_min, x2_max = x.iloc[:, 1].min(), x.iloc[:, 1].max()  # 第1列的范围
t1 = np.linspace(x1_min, x1_max, N)
t2 = np.linspace(x2_min, x2_max, M)
x1, x2 = np.meshgrid(t1, t2)  # 生成网格采样点
x_show = np.stack((x1.flat, x2.flat), axis=1)  # 测试点
cm_light = mpl.colors.ListedColormap(['#A0FFA0', '#FFA0A0', '#A0C0FF'])
cm_dark = mpl.colors.ListedColormap(['r', 'r', 'b'])
y_show_hat = model.predict(x_show)  # 预测值
y_show_hat = y_show_hat.reshape(x1.shape)  # 使之与输入的形状相同
plt.figure(facecolor='w')
plt.pcolormesh(x1, x2, y_show_hat, cmap=cm_light)  # 预测值的显示
plt.scatter(x_test.iloc[:, 0], x_test.iloc[:, 1], c=y_test.ravel(), edgecolors='k', s=100, cmap=cm_dark, marker='o')  # 测试数据
# edgecolors 边缘颜色
plt.scatter(x.iloc[:, 0], x.iloc[:, 1], c=y.ravel(), edgecolors='k', s=40, cmap=cm_dark)  # 全部数据
plt.xlabel('追过人数', fontsize=15)
plt.ylabel('被追人数', fontsize=15)
plt.xlim(x1_min, x1_max)
plt.ylim(x2_min, x2_max)
plt.grid(True)
plt.title(u'单身狗数据的决策树分类-男蓝女绿', fontsize=17)
plt.show()